package com.edu.dk.security.impl;

import java.util.regex.Matcher;
import java.util.regex.Pattern;

import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;

import com.edu.dk.security.*;
import com.edu.dk.security.exception.ResponseHeaderRejectedException;
import com.edu.dk.security.util.CookieSupport;
import com.edu.dk.security.util.HumanReadableSize;
import com.edu.dk.security.util.ObjectUtils;
import com.edu.dk.security.util.StringEscapeUtils;
import com.edu.dk.security.util.StringUtils;

/**
 * 过滤header中的CRLF，将status message用HTML entities转义，限制cookie的总大小。
 * 
 */
public class ResponseHeaderSecurityFilter
		implements WebContextLifecycleInterceptor, HeaderNameInterceptor, HeaderValueInterceptor, CookieInterceptor,
		CookieHeaderValueInterceptor, StatusMessageInterceptor, RedirectLocationInterceptor {

	// Max Cookie Size Default
	public static final HumanReadableSize MAX_SET_COOKIE_SIZE_DEFAULT = new HumanReadableSize("10k");
	// CRLF正则匹配
	private static final Pattern crlf = Pattern.compile("\\r\\n|\\r|\\n");
	// Cookie大小累计器
	private final CookieLengthAccumulator cookieLengthAccumulator;
	// 实际设置的Cookie最大大小
	private HumanReadableSize maxCookieSize;

	// --------------------------------------------------------
	// Constructor
	// --------------------------------------------------------
	public ResponseHeaderSecurityFilter() {
		this(null);
	}

	public ResponseHeaderSecurityFilter(HttpServletRequest request) {
		if (request == null) {
			// 如果request为空采取ThreadLocal累计cookie的大小
			cookieLengthAccumulator = new ThreadLocalBasedCookieLengthAccumulator();
		} else {
			// 否则采取Request累计cookie的大小
			cookieLengthAccumulator = new RequestBasedCookieLengthAccumulator(request);
		}
	}

	/**
	 * 允许的cookie最大限制
	 * 
	 * @return
	 */
	public HumanReadableSize getMaxCookieSize() {
		return maxCookieSize == null || maxCookieSize.getValue() <= 0 ? MAX_SET_COOKIE_SIZE_DEFAULT : maxCookieSize;
	}

	public void setMaxCookieSize(HumanReadableSize maxSetCookieSize) {
		this.maxCookieSize = maxSetCookieSize;
	}

	@Override
	public void prepare() {
	}

	@Override
	public void commitHeaders() {
		// 清空以重新计算下次
		cookieLengthAccumulator.reset();
	}

	@Override
	public void commit() {
	}

	@Override
	public String checkHeaderName(String name) {
		// 主要是防止CRLF攻击
		if (containsCRLF(name)) {
			// Invalid response header
			String msg = "Invalid response header: " + StringEscapeUtils.escapeJava(name);
			throw new ResponseHeaderRejectedException(msg);
		}
		return name;
	}

	@Override
	public String checkHeaderValue(String name, String value) {
		// 主要是防止CRLF攻击
		return ObjectUtils.defaultIfNull(filterCRLF(value, "header " + name), value);
	}

	@Override
	public Cookie checkCookie(Cookie cookie) {
		// 主要是防止CRLF攻击
		String name = cookie.getName();
		if (containsCRLF(name)) {
			// Invalid cookie name
			return null;
		}
		String value = cookie.getValue();
		String filteredValue = filterCRLF(value, "cookie " + name);
		if (filteredValue == null) {
			return cookie;
		} else {
			CookieSupport newCookie = new CookieSupport(cookie);
			newCookie.setValue(filteredValue);
			return newCookie;
		}
	}

	@Override
	public String checkCookieHeaderValue(String name, String value, boolean setHeader) {
		// 主要是防止Cookie过大
		if (value != null) {
			int maxCookieSize = (int) getMaxCookieSize().getValue();
			int length = cookieLengthAccumulator.getLength();
			// cookie size limit
			if (length + value.length() > maxCookieSize) {
				// Cookie size exceeds the max value
				return StringUtils.EMPTY;
			} else {
				if (setHeader) {
					cookieLengthAccumulator.setCookie(value);
				} else {
					cookieLengthAccumulator.addCookie(value);
				}
			}
		}
		return value;
	}

	@Override
	public String checkStatusMessage(int sc, String msg) {
		// StatusMessage中的HTML转义
		return StringEscapeUtils.escapeHtml(msg);
	}

	@Override
	public String checkRedirectLocation(String location) {
		// 主要是防止CRLF攻击
		return ObjectUtils.defaultIfNull(filterCRLF(location, "redirectLocation"), location);
	}

	/**
	 * 是否包含CRLF
	 * 
	 * @param str
	 * @return
	 */
	private boolean containsCRLF(String str) {
		if (str != null) {
			for (int i = 0; i < str.length(); i++) {
				switch (str.charAt(i)) {
				case '\r':
				case '\n':
					return true;
				}
			}
		}
		return false;
	}

	/** 如果不包含CRLF，则返回<code>null</code>，否则除去所有CRLF，替换成空格。 */
	private String filterCRLF(String value, String logInfo) {
		if (containsCRLF(value)) {
			StringBuffer sb = new StringBuffer();
			Matcher m = crlf.matcher(value);
			while (m.find()) {
				m.appendReplacement(sb, " ");
			}
			m.appendTail(sb);
			return sb.toString();
		}
		return null;
	}

	// ------------------------------------------------------------
	// 辅助内部类
	// ------------------------------------------------------------
	/**
	 * Cookies大小计量器
	 * 
	 * @author cixao84@hotmail.com
	 *
	 */
	private static abstract class CookieLengthAccumulator {
		/**
		 * 添加Cookie时累加其大小
		 * 
		 * @param cookie
		 */
		public final void addCookie(String cookie) {
			setLength(getLength() + cookie.length());
		}

		public final void setCookie(String cookie) {
			setLength(cookie.length());
		}

		public abstract int getLength();

		protected abstract void setLength(int length);

		protected abstract void reset();
	}

	/**
	 * 基于ThreadLocal的Cookies大小计量器
	 * 
	 * @author cixao84@hotmail.com
	 *
	 */
	private final class ThreadLocalBasedCookieLengthAccumulator extends CookieLengthAccumulator {
		private final ThreadLocal<Integer> cookieLengthHolder = new ThreadLocal<Integer>();

		public int getLength() {
			Object value = cookieLengthHolder.get();

			if (value instanceof Integer) {
				return (Integer) value;
			} else {
				return 0;
			}
		}

		protected void setLength(int length) {
			cookieLengthHolder.set(length);
		}

		protected void reset() {
			cookieLengthHolder.remove();
		}
	}

	/**
	 * 基于Request存储的Cookies大小计量器
	 * 
	 * @author cixao84@hotmail.com
	 *
	 */
	private final class RequestBasedCookieLengthAccumulator extends CookieLengthAccumulator {
		private static final String COOKIE_LENGTH_ATTR = "_COOKIE_LENGTH_";// 存储cookie大小的request属性名称
		private final HttpServletRequest request;

		private RequestBasedCookieLengthAccumulator(HttpServletRequest request) {
			this.request = request;
		}

		public int getLength() {
			Object value = request.getAttribute(COOKIE_LENGTH_ATTR);

			if (value instanceof Integer) {
				return (Integer) value;
			} else {
				return 0;
			}
		}

		protected void setLength(int length) {
			request.setAttribute(COOKIE_LENGTH_ATTR, length);
		}

		protected void reset() {
			request.removeAttribute(COOKIE_LENGTH_ATTR);
		}
	}

}